import asyncio
import aiohttp
import json
from typing import List, Optional
import logging


class PaperInfo:
    def __init__(self, paperId, title, abstract, venue, year, citationCount, tldr, citationStyles, authors):
        self.paperId = paperId
        self.title = title
        self.abstract = abstract
        self.venue = venue
        self.year = year
        self.citationCount = citationCount
        # self.tIdr = tIdr
        self.citation = citationStyles['bibtex']
        # self.authors = authors

    def __str__(self):
        return f"{self.title} {self.year} {self.citationCount}"

    def __eq__(self, other):
        return self.paperId == other.paperId


class SemScholarHttpClient:
    def __init__(self, api_key: str):
        self.logger = logging.getLogger(__name__)
        # assert config_path.exists()
        # with open(config_path) as fp:
        #     self.api_key = json.load(fp)['api_key']
        self.api_key = api_key
        self.headers = {'x-api-key': self.api_key}

    async def do_request(
            self,
            req: str,
            fields='title,year,abstract,authors.name,venue,citationCount,citationStyles,tldr'
    ) -> Optional[dict]:
        query_params = {
            'query': req,
            'limit': 100,
            'year': '2016-',
            'fields': fields,
            'order': 'citationCount:desc'
        }
        args = dict(
            url='https://api.semanticscholar.org/graph/v1/paper/search',
            params=query_params,
            headers=self.headers
        )
        try:
            async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=10)) as session:
                async with session.get(**args) as response:
                    if response.status == 200:
                        return await response.json()
                    else:
                        self.logger.info(response.content)
                        return None
        except aiohttp.ClientConnectionError as e:
            self.logger.info('client connection error')
        except aiohttp.ClientResponseError as e:
            self.logger.info('client response error')
        except asyncio.exceptions.TimeoutError as e:
            self.logger.info('timeout')
        return None

    async def get_papers(self, req: str, error_sleep_time: int = 1, n_papers: Optional[int] = None) -> List[PaperInfo]:
        resp = None
        while resp is None:
            await asyncio.sleep(error_sleep_time)
            self.logger.info('base request')
            resp = await self.do_request(req=req)
        if not 'data' in resp.keys():
            return []
        papers = [PaperInfo(**paper) for paper in resp['data']]
        # for paper in papers:
        #     self.logger.info(paper)
        self.logger.info(f"{len(papers)} found")
        papers = [paper for paper in papers if paper.abstract is not None]
        self.logger.info(f"{len(papers)} with abstract")
        papers = list(sorted(papers, key=lambda x: x.citationCount))
        if n_papers is None:
            return papers
        elif isinstance(n_papers, int):
            return papers[:n_papers]
        else:
            raise Exception(f"Undefined class of n_papers arg {n_papers.__class__}")


async def _test_example():
    with open('credentials/config.json') as fp:
        api_key = json.load(fp)['api_key']
    sem_scholar_client = SemScholarHttpClient(api_key=api_key)
    papers = await sem_scholar_client.get_papers(req='general additive model')
    # papers = sem_scholar_client.get_papers(req='Объяснительные методы машинного обучения')
    # papers = sem_scholar_client.get_papers(req='Анализ выживаемости в машинном обучении')
    assert all([paper is not None for paper in papers])


if __name__ == '__main__':
    asyncio.run(_test_example())
